package com.asiainfo.dacp.common.twophasesink;



import org.apache.flink.api.common.state.ListState;
import org.apache.flink.api.common.state.ListStateDescriptor;
import org.apache.flink.api.common.typeinfo.TypeHint;
import org.apache.flink.api.common.typeinfo.TypeInformation;
import org.apache.flink.api.common.typeinfo.Types;
import org.apache.flink.configuration.Configuration;
import org.apache.flink.runtime.state.FunctionInitializationContext;
import org.apache.flink.runtime.state.FunctionSnapshotContext;
import org.apache.flink.streaming.api.environment.StreamExecutionEnvironment;
import org.apache.flink.streaming.api.functions.sink.TwoPhaseCommitSinkFunction;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;


import java.sql.Connection;
import java.sql.DriverManager;
import java.sql.SQLException;
import java.util.ArrayList;
import java.util.Date;

/**
 * Doris 二阶段提交 Sink
 * 该类实现了 Flink 的二阶段提交 Sink 函数，用于将数据批量写入 Doris 数据库，
 * 通过二阶段提交协议确保数据的一致性和精确一次语义。
 * 数据输入类型  事务上下文类型  上下文类型
 */

public class DorisBatchTwoPhaseCommitSink
        extends TwoPhaseCommitSinkFunction<String, DorisTransaction, Void> {

    private static final Logger log = LoggerFactory.getLogger(DorisBatchTwoPhaseCommitSink.class);

    private final String jdbcUrl;
    private final String username;
    private final String password;
    private final String tableName;
    private final int batchSize;

    private transient ListState<DorisTransaction> checkpointedState;

    /**
     * 构造 Doris 批量二阶段提交 Sink 实例
     *
     * @param jdbcUrl   Doris 数据库 JDBC 连接地址
     * @param username  数据库用户名
     * @param password  数据库密码
     * @param tableName 目标表名
     * @param batchSize 批处理大小
     * @param env       Flink 执行环境，用于创建序列化器
     */

    public DorisBatchTwoPhaseCommitSink(String jdbcUrl, String username,
                                        String password, String tableName, int batchSize,
                                        StreamExecutionEnvironment env) {


super(TypeInformation.of(DorisTransaction.class).createSerializer(env.getConfig()),
      Types.VOID.createSerializer(env.getConfig()));


        this.jdbcUrl = jdbcUrl;
        this.username = username;
        this.password = password;
        this.tableName = tableName;
        this.batchSize = batchSize;
    }

    /**
     * 打开 Sink 连接，初始化操作
     *
     * @param parameters 配置参数
     * @throws Exception 初始化异常
     */

    @Override
    public void open(Configuration parameters) throws Exception {
        super.open(parameters);
        // 注册 Doris JDBC 驱动
        Class.forName("com.mysql.jdbc.Driver");
    }

    /**
     * 开始事务 - 创建事务上下文
     * 建立数据库连接并设置手动提交模式
     *
     * @return DorisTransaction 事务上下文对象
     * @throws Exception 数据库连接异常
     */

    @Override
    protected DorisTransaction beginTransaction() throws Exception {
        log.info("开始新事务");
        Connection connection = DriverManager.getConnection(jdbcUrl, username, password);
        connection.setAutoCommit(false);
        return new DorisTransaction(connection, new ArrayList<>(), tableName, batchSize);
    }

    /**
     * 执行写入操作 - 批量缓存数据
     * 将数据添加到事务缓冲区，达到批处理大小时执行批量写入
     *
     * @param transaction 当前事务上下文
     * @param value       输入数据记录
     * @param context     Sink 上下文
     * @throws Exception 数据处理异常
     */

    @Override
    protected void invoke(DorisTransaction transaction, String value, Context context) throws Exception {
        transaction.addRecord(value);

        // 如果达到批处理大小，执行批量写入
        if (transaction.shouldFlush()) {
            transaction.executeBatch();
        }
    }

    /**
     * 预提交阶段 - 执行批量写入但不提交
     * 确保所有缓冲数据都写入数据库，为正式提交做准备
     *
     * @param transaction 当前事务上下文
     * @throws Exception 预提交异常
     */

    @Override
    protected void preCommit(DorisTransaction transaction) throws Exception {
        log.info("预提交事务，当前批次大小: {}", transaction.getRecordCount());
        // 确保所有数据都写入
        if (transaction.getRecordCount() > 0) {
            transaction.executeBatch();
        }
    }

    /**
     * 提交事务
     * 正式提交数据库事务，确保数据持久化
     *
     * @param transaction 当前事务上下文
     */

    @Override
    protected void commit(DorisTransaction transaction) {
        try {
            log.info("提交事务");
            transaction.commit();
        } catch (SQLException e) {
            log.error("提交事务失败", e);
            throw new RuntimeException("提交事务失败", e);
        } finally {
            transaction.close();
        }
    }

    /**
     * 中止事务
     * 回滚数据库事务，用于处理异常情况
     *
     * @param transaction 当前事务上下文
     */

    @Override
    protected void abort(DorisTransaction transaction) {
        try {
            log.warn("中止事务");
            transaction.rollback();
        } catch (SQLException e) {
            log.error("中止事务失败", e);
        } finally {
            transaction.close();
        }
    }

    /**
     * 初始化状态
     * 用于恢复检查点状态，确保故障恢复后的一致性
     *
     * @param context 函数初始化上下文
     * @throws Exception 状态初始化异常
     */

    @Override
    public void initializeState(FunctionInitializationContext context) throws Exception {
        super.initializeState(context);

        ListStateDescriptor<DorisTransaction> descriptor =
                new ListStateDescriptor<>("doris-transaction-state",
                        TypeInformation.of(new TypeHint<DorisTransaction>() {
                        }));
        checkpointedState = context.getOperatorStateStore().getListState(descriptor);
    }

    @Override
    public void snapshotState(FunctionSnapshotContext context) throws Exception {
        super.snapshotState(context);
    }
}